##########################################################################################

library(ComplexHeatmap)
library(dplyr)
library(ggplot2)
library(data.table)
library(RColorBrewer)
library(optparse)

##########################################################################################

option_list <- list(
    make_option(c("--maf_path"), type = "character") ,
    make_option(c("--images_path"), type = "character") ,
    make_option(c("--info_file"), type = "character") ,
    make_option(c("--class_order_file"), type = "character") ,
    make_option(c("--smg_list"), type = "character")
)

if(1!=1){
    
    work_dir <- "~/20220915_gastric_multiple/dna_combine"
    maf_path <- paste(work_dir,"/","maf",sep="")
	images_path <- paste(work_dir,"/","images",sep="")
	info_file <- paste(work_dir,"/config/tumor_normal.class.list",sep="")
	class_order_file <- paste(work_dir,"/config/Class_order.list",sep="")
	smg_list <- paste(work_dir,"/public_ref/SMG_sort.list",sep="")

}

###########################################################################################

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

maf_path <- opt$maf_path
info_file <- opt$info_file
images_path <- opt$images_path
class_order_file <- opt$class_order_file
smg_list <- opt$smg_list

###########################################################################################
smg <- data.frame(fread(smg_list , header = F))
colnames(smg) <- "Gene_Symbol"

class_order <- data.frame(fread(class_order_file , header = T))
info <- data.frame(fread(info_file))
info$Class <- factor(info$Class , levels = class_order$Class , ordered=T)

###########################################################################################

col <- c(
  brewer.pal(9,"YlGnBu")[6],
  rgb(234,106,79,alpha=255,maxColorValue=255),
  rgb(203,24,30,alpha=255,maxColorValue=255),
  rgb(255,0,0,alpha=255,maxColorValue=255)
  )

names(col) <- c("IM" , "IGC" , "DGC" , "GC")

Variant_Types <- c("Missense_Mutation","Nonsense_Mutation","Frame_Shift_Ins","Frame_Shift_Del","In_Frame_Ins","In_Frame_Del","Splice_Site","Nonstop_Mutation")

col_class <- col[1:3]

###########################################################################################
CombineMaf <- function(info = info , maf_path = maf_path){
	maf_all <- c()
	for(i in 1:nrow(info)){
		Tumor <- info[i , "Tumor" ]
		Normal <- info[i , "Normal"]
		Class <- info[i , "Class"]

		file <- paste0(maf_path,"/",Tumor,"_",Normal,"_GGA_Filter_funcotated.maf")
		maf <- read.csv(file,comment.char="#",sep="\t")

		## 此样本必须有突变
		maf <- maf[maf$t_alt_count>0,]

		maf_use <- data.frame(Hugo_Symbol = maf$Hugo_Symbol, 
			Chromosome = maf$Chromosome , Start_Position =  maf$Start_Position , End_Position = maf$End_Position ,
			Reference_Allele = maf$Reference_Allele , Tumor_Seq_Allele2 = maf$Tumor_Seq_Allele2 , 
			Variant_Classification = maf$Variant_Classification , 
			Tumor = Tumor , Normal = Normal , Class = Class , stringsAsFactors=FALSE )

		maf_all <- rbind(maf_all,maf_use)
	}
	return(maf_all)
}

CreateMutMatrix <- function(dat = dat , smg = smg , Variant_Type = Variant_Type ){

	pre <- c("IM")
	can <- c("IGC","DGC")

	mut <- dat
	mut$Variant_Classification <- as.character(mut$Variant_Classification)
	mut <- mut[which(mut$Variant_Classification %in% Variant_Type),]

	## INDEL 和 INS 合并
	mut[grep("In_Frame",mut$Variant_Classification),'Variant_Classification'] = "In_Frame"
	mut[grep("Frame_Shift",mut$Variant_Classification),'Variant_Classification'] = "Frame_Shift"

	## Gene只考虑有突变的SMG和CGC
	Gene <- smg[smg$Gene_Symbol %in% mut$Hugo_Symbol,]

	## 构建矩阵
	Sample <- info[,"Tumor"]
	maf_matrix <- matrix("" , ncol = length(Sample) , nrow = length(Gene) , dimnames = list(Gene,Sample))

	for(gene in rownames(maf_matrix)){
		print(gene)

		## 突变类型
		for(tumor in colnames(maf_matrix)){
			index <- which(mut$Hugo_Symbol==gene & mut$Tumor==tumor)
			if(length(index)==0){
				var=""
			}else if(length(index)==1){
				var <- mut[index,'Variant_Classification']
			}else if(length(index)>1){
				var <- "Multiple_Hits"
			}
			print(var)
			maf_matrix[ which(rownames(maf_matrix)==gene) , which(colnames(maf_matrix)==tumor) ] <-  var
		}

		## 是否为共享
		tmp <- mut[mut$Hugo_Symbol==gene,]
		tmp$Location <- apply(tmp , 1 , function(x){paste(x[c(1:6,9)], collapse =";") } )

		tmp_1 <- tmp %>% 
		group_by(Location) %>% 
		summarize( Chromosome = Chromosome[1] , Start_Position = Start_Position[1] , End_Position = End_Position[1] , 
			Reference_Allele = Reference_Allele[1] , Tumor_Seq_Allele2 = Tumor_Seq_Allele2[1] , Tumor = Tumor ,
			Share = ifelse(length(which(Class %in%  pre))!=0 & length(which(Class %in% can ))!=0 , "Share" , "" ) )
		
		tmp_1 <- data.frame(tmp_1)
		## Private的是小白框
		tmp_1 <- subset(tmp_1,Share!="Share")

		if(dim(tmp_1)[1]!=0){
			maf_matrix[rownames(maf_matrix)==gene ,  colnames(maf_matrix) %in% tmp_1$Tumor ] <- paste0(maf_matrix[rownames(maf_matrix)==gene ,  colnames(maf_matrix) %in% tmp_1$Tumor ] , ";Private" )
		}
	}

	return(maf_matrix)
}

MutMatrixOrder <- function( mut = mut , info = info){

	######################################
	## 肿瘤的驱动基因

	## 将突变的百分比放在基因名上
	## 百分比为人的
	mut_per <- paste0( 100 * round(
			apply(mut , 1 , function(x){length(unique(info[info$Tumor %in% names(which(x!="")) ,"Normal"]))} ) / 
			length(unique(info$Normal)) ,
		2)) 
	mut_per <- as.numeric(mut_per)

	rownames(mut) <- paste0(rownames(mut) , " ( " , mut_per , "%" , " ) ")

	## 计算有共享突变的百分比
	## mut_share <- paste0( 100 * round(
	##		apply(mut , 1 , function(x){length(unique(info[info$Tumor %in% names(x[grep("Share",x)]) ,"Normal"]))} ) / 
	##		apply(mut , 1 , function(x){length(unique(info[info$Tumor %in% names(which(x!="")) ,"Normal"]))} ) ,
	##	2))
	## rownames(mut) <- paste0(rownames(mut) , " ( " , mut_per , "% , " , mut_share , "%" , " ) ")

	## 基因排序
	## 按照患肿瘤的病人，然后按照癌前病变有的病人
	IGC_sample <- info[info$Class=="IGC","Tumor"]
	DGC_sample <- info[info$Class=="DGC","Tumor"]
	IM_sample <- info[info$Class=="IM","Tumor"]

	mut_tumor <- mut[,colnames(mut) %in% c(IGC_sample,DGC_sample)]
	mut_pre <- mut[,colnames(mut) %in% c(IM_sample)]

	NumMut_Tumor <- apply(mut_tumor , 1 ,function(x){
		length(unique(info[info$Tumor %in% names(which(x!="")) ,"Normal"]))}
	)

	NumMut_Pre <- apply(mut_pre , 1 ,function(x){
		length(unique(info[info$Tumor %in% names(which(x!="")) ,"Normal"]))}
	)

	MutNumMatrix <- data.frame(cbind(NumMut_Pre,NumMut_Tumor))

	## 基因排序
	mut <- mut[order(MutNumMatrix$NumMut_Tumor , MutNumMatrix$NumMut_Pre , decreasing = T ),]

	## 样本排序
	## 保证前面的mut的样本顺序和info的一致
	mut <- rbind( mut , Class = info$Class )
	mut <- mut[ , order(mut["Class",] , 
		mut[1,] , mut[2,] , mut[3,] , mut[4,] , mut[5,] , mut[6,] , 
		mut[7,] , mut[8,] , mut[9,] , mut[10,] , mut[11,] , mut[12,] ,
		mut[13,] , mut[14,] , mut[15,] , mut[16,] , mut[17,] , mut[18,] ,
		colnames(mut) , decreasing = T )]

	## 去除Class
	mut <- mut[rownames(mut)!=c("Class"),]

	return(mut)
}

SumSilentAndNonSilent <- function(sample_order = sample_order , info = info , maf_path = maf_path , Variant_Type = Variant_Type ){
	MutNum_all <- c()
	for(i in 1:length(sample_order)){
		print(i)

		Tumor <- sample_order[i]
		Normal <- info[info$Tumor==Tumor,"Normal"]
		
		file <- paste0(maf_path,"/",Tumor,"_",Normal,"_GGA_Filter_funcotated.maf")
		maf <- read.csv(file,comment.char="#",sep="\t")
		
		## 突变必须存在
		maf <- maf[maf$t_alt_count>0,]


		Silent_num <- length(which(maf$Variant_Classification %in% "Silent"))
		NonSilent_num <- length(which(maf$Variant_Classification %in% Variant_Type))

		MutNum_all <- rbind( MutNum_all , data.frame(NonSilent_num = NonSilent_num , Silent_num = Silent_num  ))
	}

	return(MutNum_all)

}

plotMutWaterFull <- function( MutMatrix = MutMatrix , col_class = col_class , Variant_Type_Combine = Variant_Type_Combine , images_name = images_name ){

	################################################################################################
	## 注释的名字
	annotation_name <- c("Class","Mut_Num")

	## 展示出现频率 >= 5个样本的突变
	mut <- MutMatrix[apply(MutMatrix,1,function(x){length(which(x!=""))>5}),]

	## 突变矩阵排序
	mut <- MutMatrixOrder( mut = mut , info = info )
	
	## 顺序,很重要！
	class_order <- info[order(info$Class , decreasing=T),"Class"]
	sample_order <- colnames(mut)

	################################################################################################
	## 突变的颜色
	col = c(rgb(red=48,green=115,blue=186,alpha=255,max=255),
		rgb(red=236,green=27,blue=35,alpha=255,max=255),
		rgb(red=236,green=179,blue=33,alpha=255,max=255),
		rgb(red=235,green=230,blue=26,alpha=255,max=255),
		rgb(red=150,green=131,blue=189,alpha=255,max=255),
		rgb(red=65,green=174,blue=119,alpha=255,max=255),
		"white")

	names(col) = c(
	  'Missense_Mutation',
	  'Nonsense_Mutation',
	  'Frame_Shift',
	  'In_Frame',
	  'Splice_Site',
	  'Multiple_Hits',
	  'Private'
	)

	## 设置不同背景的颜色
	alter_fun = list(
	    background = function(x, y, w, h) {
	        grid.rect(x, y, w-unit(0.5, "mm"), h-unit(0.5, "mm"), 
	            gp = gpar(fill = "#F2F2F2", col = NA))
	    },
	    Missense_Mutation = function(x, y, w, h) {
	        grid.rect(x, y, w-unit(0.5, "mm"), h-unit(0.5, "mm"), 
	            gp = gpar(fill = col["Missense_Mutation"], col = NA))
	    },
	    Nonsense_Mutation = function(x, y, w, h) {
	        grid.rect(x, y, w-unit(0.5, "mm"), h-unit(0.5, "mm"), 
	            gp = gpar(fill = col["Nonsense_Mutation"], col = NA))
	    },
	  
	    Nonstop_Mutation = function(x, y, w, h) {
	        grid.rect(x, y, w-unit(0.5, "mm"), h-unit(0.5, "mm"),
	            gp = gpar(fill = col["Nonstop_Mutation"], col = NA))
	    },
	    Frame_Shift = function(x, y, w, h) {
	        grid.rect(x, y, w-unit(0.5, "mm"), h-unit(0.5, "mm"),
	            gp = gpar(fill = col["Frame_Shift"], col = NA))
	    },
	    In_Frame = function(x, y, w, h) {
	        grid.rect(x, y, w-unit(0.5, "mm"), h-unit(0.5, "mm"),
	            gp = gpar(fill = col["In_Frame"], col = NA))
	    },
	    Splice_Site = function(x, y, w, h) {
	        grid.rect(x, y, w-unit(0.5, "mm"), h-unit(0.5, "mm"),
	            gp = gpar(fill = col["Splice_Site"], col = NA))
	    },
	    Multiple_Hits = function(x, y, w, h) {
	        grid.rect(x, y, w-unit(0.5, "mm"), h-unit(0.5, "mm"),
	            gp = gpar(fill = col["Multiple_Hits"], col = NA))
	    },
	    Private = function(x, y, w, h) {
	        grid.rect(x, y, w*0.5, h*0.4,,
	            gp = gpar(fill = col["Private"], col = NA))
	    },
	    show_legend = FALSE 
	)

	################################################################################################
	## 顶部注释 突变数量
	MutSum <- SumSilentAndNonSilent(sample_order = sample_order , info = info , maf_path = maf_path , Variant_Type = Variant_Type )

	col_Silent <- c(
		rgb(red=212,green=216,blue=236,alpha=255,max=255), ##NonSilent
		rgb(red=40,green=52,blue=146,alpha=255,max=255) ## Silent
	)
	names(col_Silent) <- c("NonSilent_num","Silent_num")


	## 顶部注释 class
	top_annotation <- HeatmapAnnotation(
		foo = anno_empty(height = unit(2, "cm") , border =F) ,  ## 留给注释空间
		Mut_Num = anno_barplot(
			MutSum, 
    		gp = gpar(fill = col_Silent, col = col_Silent),
    		border = FALSE,
    		height = unit(2.5, "cm")
    	) ,
		Class = class_order ,
		col = list( 
			Class = col_class ,
			Mut_Num = col_Silent
		) ,

		annotation_name_side = "left" , 
	  	border = T ,
	  	gap = unit(1, "mm") ,
	  	show_annotation_name = c(Mut_Num = FALSE) , 
	  	annotation_name_gp = gpar(fontsize = 12),

	  	show_legend = FALSE  ## 所有的legend均后期定义
	)

	################################################################################################
	## 左部加基因在样本中分布的柱状图,癌前和癌分开
	## 计算肿瘤中和癌前的突变的病人数
	IGC_sample <- info[info$Class=="IGC","Tumor"]
	DGC_sample <- info[info$Class=="DGC","Tumor"]
	IM_sample <- info[info$Class=="IM","Tumor"]

	mut_tumor <- mut[,colnames(mut) %in% c(IGC_sample,DGC_sample)]
	mut_pre <- mut[,colnames(mut) %in% c(IM_sample)]

	NumMut_Tumor <- apply(mut_tumor , 1 ,function(x){
		length(unique(info[info$Tumor %in% names(which(x!="")) ,"Normal"]))}
	)

	NumMut_Pre <- apply(mut_pre , 1 ,function(x){
		length(unique(info[info$Tumor %in% names(which(x!="")) ,"Normal"]))}
	)

	MutNumMatrix <- cbind(NumMut_Pre,NumMut_Tumor)

	extend <- max(NumMut_Tumor) / max(NumMut_Pre) -1

	sample_num <- length(unique(info$Normal))

	left_annotation = rowAnnotation(
		Tumor = anno_barplot(MutNumMatrix[,"NumMut_Tumor"],
			border = FALSE ,
			gp = gpar(fill = rgb(red=204,green=148,blue=98,alpha=255,max=255) ),
			axis_param = list(direction = "reverse" , labels_rot = 0 , side = "bottom"),
			bar_width = 1, width = unit(3, "cm") ,
		),
		Tumor_Num = anno_text(
			paste0(round(as.numeric(MutNumMatrix[,"NumMut_Tumor"]/sample_num) * 100 ) , "%"), 
			location = unit(3, "mm"), just = "right"
	    ),
		PreCancerous = anno_barplot(
			MutNumMatrix[,"NumMut_Pre"], 
			border = FALSE ,
			gp = gpar(fill = rgb(red=196,green=60,blue=132,alpha=255,max=255) ),
			axis_param = list(labels_rot = 0 , side = "bottom"),
   			bar_width = 1, width = unit(3, "cm") ,
   			extend = extend ## 让Tumor和Pre坐标轴一致
   		),
   		Pre_Num = anno_text(
   			paste0(round(as.numeric(MutNumMatrix[,"NumMut_Pre"]/sample_num) * 100 ) , "%"), 
   			location = unit(-38, "mm") , just = "right"
	    ),
   		gap = unit(7, "mm"), ## 两注释分开来
   		annotation_name_side = "top",
   		annotation_name_gp = gpar(fontsize = 10),
   		annotation_name_rot = 0,
   		annotation_name_offset = unit(2, "mm")
   	) 

	################################################################################################
	## 右部加总的突变的样本数，还有癌前和癌共享的样本
	## 计算肿瘤中和癌前的突变的病人数
	mut_per <- as.numeric(apply(mut , 1 , function(x){length(unique(info[info$Tumor %in% names(which(x!="")) ,"Normal"]))} ))
			
	## 计算有共享突变的百分比
	mut_share <-  as.numeric(apply(mut , 1 , function(x){length(unique(info[info$Tumor %in% names(which(x[grep("Private",x,invert=T) ]!="")) ,"Normal"]))} ) )

	mutSum_share <- cbind(Share = mut_share ,Total = mut_per - mut_share )

	col_Share <- c(
		rgb(red=33,green=113,blue=181,alpha=255,max=255), ## Share
		rgb(red=204,green=204,blue=204,alpha=255,max=255)## No Share
		
	)
	names(col_Share) <- c("Share","Total")

	## 右注释 class
	right_annotation <- rowAnnotation(
		mutSum_share = anno_barplot(
			mutSum_share, 
    		gp = gpar(fill = col_Share, col = col_Share),
    		border = FALSE,
    		height = unit(2.5, "cm") ,
    		width = unit(3, "cm") ,
    	) ,
		col = list( 
			mutSum_share = col_Share
		) ,

		annotation_name_side = "top",
	  	border = T ,
	  	gap = unit(1, "mm") ,
	  	show_annotation_name = c(mutSum_share = FALSE) , 
	  	annotation_name_gp = gpar(fontsize = 12),
	  	show_legend = FALSE  ## 所有的legend均后期定义
	)

	## 给沉默突变和非沉默突变加Legend
	lgd_Mut_Num = Legend(labels = c("Non-Silent", "Silent","Share"), title = "MutType",
    	legend_gp = gpar(fill = c(col_Silent,col_Share[1]) , fontsize = 12) , ncol = 3 )

	ldg_Class = Legend(labels = names(col_class) , title = "Pathogenic" ,
	 legend_gp = gpar(fill = col_class , bar_width = 1 , fontsize = 12) , ncol = 3 , 
	 gap = unit(1, "cm") ## 图例的间隔
	)

	## 突变类型注释
	col_Mut <- col 
	names(col_Mut)[length(col_Mut)] <- "Private" 

	ldg_Variant = Legend(labels = names(col_Mut) , title = "Mutations" , border = "black" , ## 注释边框加黑
	 	legend_gp = gpar(fill = col_Mut , bar_width = 1 ,fontsize = 12) , ncol = 3 , 
	 	gap = unit(1, "cm") ## 图例的间隔
	)	

	lgd_all <- packLegend(lgd_Mut_Num  , ldg_Class , ldg_Variant , column_gap = unit(1, "cm") , row_gap = unit(10, "mm") , direction = "horizontal"  )

	################################################################################################
	p <- oncoPrint(mut, name = "cases", ## 后面加分割线用的name
	    alter_fun = alter_fun, col = col, 
	    top_annotation = top_annotation  ,
	    row_names_side = "left", row_names_gp = gpar(fontsize = 20) ,  ## 基因名移到左边
	    left_annotation = left_annotation , right_annotation = right_annotation ,  ## 基因在样本间的柱状分布图
	    show_pct = FALSE , ##不展示百分比
	    border = TRUE,
	    row_order = 1:nrow(mut) ,
	    column_order = sample_order,
	    show_heatmap_legend = FALSE
	)

	##
	pdf(images_name , width = 26 , height = 9)
	draw(p )
	#draw(ldg_Variant, x = unit(0.3, "npc"), y = unit(0.005, "npc"), just = c("left", "bottom"))
	draw(lgd_all, x = unit(0.3, "npc"), y = unit(0.99, "npc"), just = c("left", "top"))
	# draw(ldg_Variant, x = unit(0.99, "npc"), y = unit(0.99, "npc"), just = c("right", "top"))

	################################################################################################
	#### 癌前和癌加分割线
	## 主体
	decorate_heatmap_body("cases", {
	    i =  which(colnames(mut) %in% IM_sample )[1] -1 
	    x = i/ncol(mut)
	    grid.lines(c(x, x), c(0, 1), gp = gpar(lwd = 0.8, lty = 2)) ## lwd代表粗细，lty代表线的类型
	})

	## 注释在对应位置加分割线
	for(anno in annotation_name){
		decorate_annotation(c(anno ), {
		    i =  which(colnames(mut) %in% IM_sample )[1] -1 
		    x = i/ncol(mut)
		    grid.lines(c(x, x), c(0, 1), gp = gpar(lwd = 0.8, lty = 2)) ## lwd代表粗细，lty代表线的类型
		})
	}

	dev.off()
}


###########################################################################################
## 合并所有的MAF文件
dat <- CombineMaf(info = info , maf_path = maf_path )

## 产生输入矩阵
Variant_Type <- c("Missense_Mutation","Nonsense_Mutation","Frame_Shift_Ins","Frame_Shift_Del","In_Frame_Ins","In_Frame_Del","Splice_Site","Nonstop_Mutation")
Variant_Type_Combine <- c("Missense_Mutation","Nonsense_Mutation","Frame_Shift","In_Frame","Splice_Site","Nonstop_Mutation","Multiple_Hits")

MutMatrix <- CreateMutMatrix( dat = dat , smg = smg , Variant_Type = Variant_Type )

## 瀑布图
images_name <- paste0(images_path , "/Mut_WaterFall.reportSMG.pdf")
plotMutWaterFull( MutMatrix = MutMatrix , col_class = col_class , Variant_Type_Combine = Variant_Type_Combine , images_name = images_name )


